// Copyright (c) 2021 Presto Labs Pte. Ltd.
// Author: chensili

#include "coin2/strategy/system.pb.h"
#include "coin2/strategy/util/reporter2.h"

#include <gtest/gtest.h>

namespace coin2::strategy::util {

class StrategyReporter2Test : public ::testing::Test {
 protected:
  StrategyReporter2Test() {
    auto strat_config = coin2::strategy::StrategyManagerConfig();
    strat_config.set_strategy_group("dmm");
    strat_config.set_strategy_name("dmm_okex_spot_pilot");
    strat_reporter_.reset(new StrategyReporter2(strat_config));
  }

  void GetAggPnlBalance(std::unordered_map<std::string, double>* agg_pnl_bal_map) {
    strat_reporter_->GetAggPnlBalance(agg_pnl_bal_map);
  }

  std::unique_ptr<StrategyReporter2> strat_reporter_;
};

TEST_F(StrategyReporter2Test, RegisterExecutorConfig) {
  // invalid acct
  auto acct_req = AccountRequestProto();
  auto mea = MarketExchangeApi::FromString("Spot.Binance.v1");
  auto product = CreateProductFromNormString(
      mea, "BTC-USDT", GetCurrentTimestamp());
  auto executor_config = PassiveExecutorConfig();
  EXPECT_ANY_THROW(strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req, *product, executor_config));

  // acct & product mismatch
  acct_req.set_market_type("Spot");
  acct_req.set_exchange("Binance");
  acct_req.set_owner("xunke00");
  mea = MarketExchangeApi::FromString("Futures.Binance.v1");
  product = CreateProductFromNormString(mea, "BTC-USD.QUARTER", GetCurrentTimestamp());
  EXPECT_ANY_THROW(strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req, *product, executor_config));

  // invalid executor config
  mea = MarketExchangeApi::FromString("Spot.Binance.v1");
  product = CreateProductFromNormString(mea, "BTC-USDT", GetCurrentTimestamp());
  EXPECT_ANY_THROW(strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req, *product, executor_config));

  // multi spot quote
  executor_config.set_lot_size(0.5);
  executor_config.set_min_pos(1);
  executor_config.set_max_pos(2);
  strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req, *product, executor_config);
  product = CreateProductFromNormString(
      mea, "ETH-BTC", GetCurrentTimestamp());
  EXPECT_ANY_THROW(strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req, *product, executor_config));

  // correct one
  product = CreateProductFromNormString(
      mea, "BTC-USDT", GetCurrentTimestamp());
  EXPECT_NO_THROW(strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req, *product, executor_config));
  product = CreateProductFromNormString(
      mea, "ETH-USDT", GetCurrentTimestamp());
  EXPECT_NO_THROW(strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req, *product, executor_config));
}

TEST_F(StrategyReporter2Test, onAccountInfo) {
  using coin2::exchange::base::order::AccountInfo;
  using exchange::base::order::OrderUpdate;

  // unregistered account
  auto mea0 = MarketExchangeApi::FromString("Spot.Binance.v1");
  AccountInfo account_info0(mea0);
  account_info0.set_owner("xunke00");
  OrderUpdate order_update("", mea0);
  order_update.SetAccountInfo(&account_info0);
  EXPECT_NO_THROW(strat_reporter_->onAccountInfo(order_update));

  // single account
  auto mea1 = MarketExchangeApi::FromString("Spot.Binance.v1");
  std::unordered_map<std::string, double> agg_pnl_bal_map;
  CurrencyBalance cbal;
  auto acct_req1 = AccountRequestProto();
  acct_req1.set_market_type("Spot");
  acct_req1.set_exchange("Binance");
  acct_req1.set_owner("xunke00");
  auto product11 = CreateProductFromNormString(
      mea1, "BTC-USDT", GetCurrentTimestamp());
  auto executor_config11 = PassiveExecutorConfig();
  executor_config11.set_lot_size(0.5);
  executor_config11.set_min_pos(1);
  executor_config11.set_max_pos(2);
  auto product12 = CreateProductFromNormString(
      mea1, "ETH-USDT", GetCurrentTimestamp());
  auto executor_config12 = PassiveExecutorConfig();
  executor_config12.set_lot_size(5);
  executor_config12.set_min_pos(10);
  executor_config12.set_max_pos(20);

  strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req1, *product11, executor_config11);
  strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req1, *product12, executor_config12);
  strat_reporter_->UpdateMidp(*product11, std::optional<double>(30000.55));

  AccountInfo account_info1(mea1);
  account_info1.set_owner("xunke00");
  cbal.set_currency("BTC");
  cbal.set_total(1.7);
  account_info1.SetBalance("BTC", cbal);
  cbal.set_currency("ETH");
  cbal.set_total(18);
  account_info1.SetBalance("ETH", cbal);
  cbal.set_currency("USDT");
  cbal.set_total(10000);
  account_info1.SetBalance("USDT", cbal);
  OrderUpdate order_update1("", mea1);
  order_update1.SetAccountInfo(&account_info1);
  strat_reporter_->onAccountInfo(order_update1);
  GetAggPnlBalance(&agg_pnl_bal_map);
  EXPECT_EQ(agg_pnl_bal_map.size(), 0);

  strat_reporter_->UpdateMidp(*product12, std::optional<double>(2500.25));
  GetAggPnlBalance(&agg_pnl_bal_map);
  EXPECT_EQ(agg_pnl_bal_map.size(), 1);
  EXPECT_DOUBLE_EQ(agg_pnl_bal_map["USDT"], 23500.86);

  // multi accounts
  auto mea2 = MarketExchangeApi::FromString("Spot.Huobi.v1");
  auto acct_req2 = AccountRequestProto();
  acct_req2.set_market_type("Spot");
  acct_req2.set_exchange("Huobi");
  acct_req2.set_owner("xunke00");
  auto product21 = CreateProductFromNormString(
      mea2, "EOS-USDT", GetCurrentTimestamp());
  auto executor_config21 = PassiveExecutorConfig();
  executor_config21.set_lot_size(5000);
  executor_config21.set_min_pos(10000);
  executor_config21.set_max_pos(20000);
  strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req2, *product21, executor_config21);

  AccountInfo account_info2(mea2);
  account_info2.set_owner("xunke00");
  cbal.set_currency("EOS");
  cbal.set_total(16000);
  account_info2.SetBalance("EOS", cbal);
  cbal.set_currency("ETH");
  cbal.set_total(5);
  account_info2.SetBalance("ETH", cbal);
  cbal.set_currency("USDT");
  cbal.set_total(5000);
  account_info2.SetBalance("USDT", cbal);
  OrderUpdate order_update2("", mea2);
  order_update2.SetAccountInfo(&account_info2);
  strat_reporter_->onAccountInfo(order_update2);
  GetAggPnlBalance(&agg_pnl_bal_map);
  EXPECT_EQ(agg_pnl_bal_map.size(), 0);

  strat_reporter_->UpdateMidp(*product21, std::optional<double>(3.45));
  GetAggPnlBalance(&agg_pnl_bal_map);
  EXPECT_DOUBLE_EQ(agg_pnl_bal_map["USDT"], 31950.86);

  auto mea3 = MarketExchangeApi::FromString("Futures.Huobi.v1");
  auto acct_req3 = AccountRequestProto();
  acct_req3.set_market_type("Futures");
  acct_req3.set_exchange("Huobi");
  acct_req3.set_owner("xunke00");
  auto product31 = CreateProductFromNormString(
      mea3, "BTC-USDT.QUARTER", GetCurrentTimestamp());
  auto product32 = CreateProductFromNormString(
      mea3, "BTC-USD.QUARTER", GetCurrentTimestamp());
  auto executor_config31 = PassiveExecutorConfig();
  executor_config31.set_lot_size(0.5);
  executor_config31.set_min_pos(-0.5);
  executor_config31.set_max_pos(0.5);
  strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req3, *product31, executor_config31);
  GetAggPnlBalance(&agg_pnl_bal_map);
  EXPECT_EQ(agg_pnl_bal_map.size(), 0);

  AccountInfo account_info3(mea3);
  account_info3.set_owner("xunke00");
  cbal.set_currency("USDT");
  cbal.set_total(3000.47);
  account_info3.SetBalance("USDT", cbal);
  OrderUpdate order_update3("", mea3);
  order_update3.SetAccountInfo(&account_info3);
  strat_reporter_->onAccountInfo(order_update3);
  GetAggPnlBalance(&agg_pnl_bal_map);
  EXPECT_EQ(agg_pnl_bal_map.size(), 1);

  cbal.set_currency("BTC");
  cbal.set_total(2.5);
  account_info3.SetBalance("USDT", cbal);
  strat_reporter_->onAccountInfo(order_update3);
  GetAggPnlBalance(&agg_pnl_bal_map);
  EXPECT_EQ(agg_pnl_bal_map.size(), 2);
  EXPECT_DOUBLE_EQ(agg_pnl_bal_map["USDT"], 34951.33);
  EXPECT_DOUBLE_EQ(agg_pnl_bal_map["BTC"], 2.5);

  auto mea4 = MarketExchangeApi::FromString("Spot.Okex.v1");
  auto acct_req4 = AccountRequestProto();
  acct_req4.set_market_type("Spot");
  acct_req4.set_exchange("Okex");
  acct_req4.set_owner("xunke00");
  auto product41 = CreateProductFromNormString(
      mea4, "BCHN-USDT", GetCurrentTimestamp());
  auto executor_config41 = PassiveExecutorConfig();
  executor_config41.set_lot_size(100);
  executor_config41.set_min_pos(100);
  executor_config41.set_max_pos(200);
  strat_reporter_->RegisterExecutorConfig(
      GetCurrentTimestamp(), acct_req4, *product41, executor_config41);
  strat_reporter_->UpdateMidp(*product41, std::optional<double>(400.55));
  GetAggPnlBalance(&agg_pnl_bal_map);
  EXPECT_EQ(agg_pnl_bal_map.size(), 0);
}

}  // namespace coin2::strategy::util
